package me.winter.imgcmps;

import java.io.BufferedOutputStream;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.time.OffsetDateTime;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Set;
import java.util.logging.Level;
import java.util.logging.Logger;

public class AdaptiveHuff {
	
	private int[] in;
	private int offset = 0;
	
	private TreeNode nytNode;
	private TreeNode treeRoot;
	// symbols has been seen, or appeared
	private Set<Integer> symbolSeen;
	private ArrayList<TreeNode> nodeList;
	
	private final int INTERNAL_NODE_VALUE = 512;
	private final int NYT_NODE_VALUE = 513;
	private final int PSEUDO_EOF = 511;
	
    public AdaptiveHuff() {
		symbolSeen = new HashSet<Integer>();
		nodeList = new ArrayList<TreeNode>();
		
		// initialize the nytNode;
		nytNode = new TreeNode(NYT_NODE_VALUE, 0);
		nytNode.parent = null;
		treeRoot = nytNode;
		nodeList.add(nytNode);
    }
        
	public AdaptiveHuff(int[] input) {
		
		int min = 256;
        for (int i = 0; i < input.length; ++i) {
        	if (input[i] < min)
        		min = input[i];
        }
        
        if (min < 0) {
	        offset = -min;
	        in = new int[input.length];
	        for (int i = 0; i < input.length; ++i) {
	        	in[i] = input[i] + offset;
	        }
        } else {
        	in = input.clone();
        }
        
        System.out.println("Offset = " + offset);
        
        
        symbolSeen = new HashSet<Integer>();
        nodeList = new ArrayList<TreeNode>();

        // initialize the nytNode;
        nytNode = new TreeNode(NYT_NODE_VALUE, 0);
        nytNode.parent = null;
        treeRoot = nytNode;
        nodeList.add(nytNode);
	}
	
	public void compress(String outfilename){
		StringBuilder codesStream = new StringBuilder();
		for ( int i=0; i< in.length; ++i) {
			int symbol = in[i];
			codesStream.append((getCode(symbol)));
			updateTree(symbol);
		}
                
        codesStream.append(getCode(PSEUDO_EOF));
        updateTree(PSEUDO_EOF);                
                
		
//		System.out.println("The compressed content is");
//		System.out.println(codesStream.toString());
		System.out.println(codesStream.toString().length() + " bits");
		
		String prefix = "AllImages/Compression/compressed/";
//        String prefix = "test/";
		String compressedFileName = prefix + outfilename;
		File outfile = new File(compressedFileName + ".hf");
		
		// TODO: write codes to BitOutputStream
		BitOutputStream bos = null;
		String codes = codesStream.toString();
		try {
			bos = new BitOutputStream(
					new BufferedOutputStream(
							new FileOutputStream(outfile)));
			bos.writeBits(8, offset);
			for (int i = 0; i < codes.length(); ++i)
				bos.writeBits(1, codes.charAt(i) - '0');
			bos.flush();
		} catch (FileNotFoundException e1) {
			e1.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		} finally {
			bos.close();
		}
	}
	
	public int[] uncompress(String infilename) {
        ArrayList<Integer> out = new ArrayList<Integer>();
	
        String prefix = "AllImages/Compression/compressed/";
//        String prefix = "test/";
        String compressedFileName = prefix + infilename;
        File infile = new File(compressedFileName + ".hf");

        BitInputStream bitIs = new BitInputStream(infile);
        try {
			offset = bitIs.readBits(8);
			System.out.println("Offset is " + offset);
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} 
        // TODO: find the first symbol and update the tree
        int symbol = getSymbol(bitIs);
        out.add(symbol);
        updateTree(symbol);


        // TODO: iterate the rest of BitInputStream
        // Zero goes to zero-branch, and one goes to
        // one-branch
        int bit;
        TreeNode p = treeRoot;
        try {
            while ((bit = bitIs.readBits(1)) != -1) {
                if (bit == 0) {
                    p = p.zero;
                } else {
                    p = p.one;
                }
                
                symbol = p.symbol;
                if (symbol != INTERNAL_NODE_VALUE) {
                    
	                if (symbol == NYT_NODE_VALUE) {
	                    symbol = getSymbol(bitIs);
	                    if (symbol == PSEUDO_EOF)
	                        break;
	                }
	                out.add(symbol);
	                updateTree(symbol);
	                p = treeRoot;
                }
            }
        } catch (IOException ex) {
            Logger.getLogger(AdaptiveHuff.class.getName()).log(Level.SEVERE, null, ex);
        }

        Integer[] op = (Integer [])out.toArray(new Integer[0]);
        
        int[] output = new int[op.length];
        for (int i = 0; i < op.length; ++i)
        	output[i] = op[i] - offset;
        
        return output;
	}

	/**
	 * The key component of this algorithm
	 * @param symbol
	 */
	private void updateTree(int symbol) {
		TreeNode toIncrement = null;
		
		if (! symbolSeen.contains(symbol)) {
			// TODO: create two new nodes: innerNode and newNode in which the
			// symbol contained. Then make the parent-child links between
			// innerNode, newNode and nytNode
			TreeNode innerNode = new TreeNode(INTERNAL_NODE_VALUE, 1);
			TreeNode newNode = new TreeNode(symbol, 1);
			
			innerNode.zero = nytNode;
			innerNode.one = newNode;
			innerNode.parent = nytNode.parent;
			if (nytNode.parent != null)
				nytNode.parent.zero = innerNode;
			else {
				treeRoot = innerNode;
			}
			nytNode.parent = innerNode;
			newNode.parent = innerNode;

			nodeList.add(1, innerNode);
			nodeList.add(1, newNode);
			symbolSeen.add(symbol);
			toIncrement = innerNode.parent;
		} else {
			// TODO: find the node contains the symbol
			toIncrement = findNode(symbol);
		}
		
		// TODO: increment the weight of parent node until 
		// reaches the root node
		while ( toIncrement != null ) {
			TreeNode bigNode = findHigherOrderNode(toIncrement.weight);
			if ( toIncrement != bigNode && toIncrement.parent != bigNode
					&& bigNode.parent != toIncrement)
				swapNode(toIncrement, bigNode);
			toIncrement.weight++;
			toIncrement = toIncrement.parent;
		}
		
	}
	
	private TreeNode findHigherOrderNode(int weight) {
		TreeNode temp = null;
		for (int i = nodeList.size() - 1; i >= 0; --i) {
			temp = nodeList.get(i);
			if ( temp.weight == weight )
				break;
		}
		return temp;
	}

	private String getCode(int symbol) {
		
		String codeword = "";
		
		if (symbolSeen.contains(symbol)) {
			// TODO: traverse to the symbol to obtain the codeword
			codeword = genCodeword(treeRoot, symbol, codeword);
		} else {
			// TODO: traverse to the Nyt node to get its codeword
			// and the fixed code of the symbol
			codeword = genCodeword(treeRoot, NYT_NODE_VALUE, "");
			codeword += getStaticCode(symbol);
		}
		return codeword;
	}
	
	/**
	 * Convert decimal value to 9-bit binary value
	 * @param decimal a integer to be converted
	 * @return string with 0,1
	 */
	public String getStaticCode(int decimal){
		String result = "";
		for (int i = 0; i < 9 ; ++i) {
			if ( (decimal % 2) == 0 )
				result = "0" + result;
			else 
				result = "1" + result;
			decimal /= 2;
		}
		return result;
	}

	private String genCodeword(TreeNode treeRoot, int symbol, String codeword) {
		
		if (treeRoot == null)
			return "";
		
		String cw = "";
		if (treeRoot.symbol == symbol)
			return codeword;
		else {
			cw = genCodeword(treeRoot.zero, symbol, codeword + "0");
			if (cw == "")
				return genCodeword(treeRoot.one, symbol, codeword + "1");
			else {
				return cw;
			}
		}
	}

	//Find the existing node in the tree
	private TreeNode findNode(int symbol) {
		// TODO Auto-generated method stub
		TreeNode tempNode = null;
		for (int i=0; i < nodeList.size(); ++i) {
			tempNode = nodeList.get(i);
			if ( (tempNode.symbol != INTERNAL_NODE_VALUE) &&
					(tempNode.symbol == symbol) )
				return tempNode;
		}
		return null;
	}
	
	private void swapNode(TreeNode n1, TreeNode n2) {
		//note that n1<n2
		//Swap the position in the list firstly
		int i1 = nodeList.indexOf(n1);
		int i2 = nodeList.indexOf(n2);
		nodeList.remove(n1);
		nodeList.remove(n2);
		nodeList.add( i1, n2);
		nodeList.add( i2, n1);
		
		//Swap the position in the tree then
		TreeNode p1 = n1.parent;
		TreeNode p2 = n2.parent;
		//If the two nodes have different parent node.
		if ( p1!=p2 ) {
			if ( p1.zero==n1 ) {
				p1.zero = n2;
			} else {
				p1.one = n2;
			}

			if ( p2.zero==n2 ) {
				p2.zero = n1;
			} else {
				p2.one = n1;
			}
		} else {
			p1.zero = n2;
			p1.one = n1;
		}
		n1.parent = p2;
		n2.parent = p1;
	
	}

    private int getSymbol(BitInputStream bitIs) {
        int retVal = 0;
        try {
            retVal = bitIs.readBits(9);
        } catch (IOException ex) {
            Logger.getLogger(AdaptiveHuff.class.getName()).log(Level.SEVERE, null, ex);
        }
        
        return retVal;
    }

	



	private class TreeNode implements Comparable<TreeNode> {
		
		public int symbol; // range from -255 to 255
		public int weight;
		public TreeNode zero;
		public TreeNode one;
		public TreeNode parent;
		
		public TreeNode(int symbol, int weight) {
			this.symbol = symbol;
			this.weight = weight;
		}

		@Override
		public int compareTo(TreeNode rhs) {
			// TODO Auto-generated method stub
			return this.weight - rhs.weight;
		}		
	}
	
//	public static void main(String[] args) {
//		int[] data = {0, 1, -250, 0, 250, 20, 20, 20, 30, 40, 50, 10, 20, 20, 20, 20};
//		
//		AdaptiveHuff aHuff = new AdaptiveHuff(data);
//		aHuff.compress("test");
//                
//        AdaptiveHuff bHuff = new AdaptiveHuff();
//        int[] uncompressed = bHuff.uncompress("test");
//        
//        System.out.println("The uncompressed content is");
//        for (int i : uncompressed) {
//            System.out.print(i + " ");
//        }
//        System.out.println();
//	}
}
